import asyncio

from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.base import TaskResult
from autogen_agentchat.conditions import ExternalTermination, TextMentionTermination
from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_agentchat.ui import Console
from autogen_core import CancellationToken

from autogen_ext.models.ollama import OllamaChatCompletionClient
import asyncio
# Assuming your Ollama server is running locally on port 11434.
model_client = OllamaChatCompletionClient(
        model="hhao/qwen2.5-coder-tools:latest",
        #model="qwen2.5-coder:latest",
        #model="GandalfBaum/deepseek_r1-claude3.7:latest",
        #model = "GandalfBaum/llama3.2-claude3.7:latest",
        host="http://192.168.99.142:11434", 
        model_info={
        "vision": False,
        "function_calling": True,
        "json_output": True,
        "family": "unknow",
        "structed_output":True,
        },
    )
# Create the primary agent.
primary_agent = AssistantAgent(
    "primary",
    model_client=model_client,
    system_message="You are a helpful AI assistant.",
)

# Create the critic agent.
critic_agent = AssistantAgent(
    "critic",
    model_client=model_client,
    system_message="Provide constructive feedback. Respond with 'APPROVE' to when your feedbacks are addressed.",
)

# Define a termination condition that stops the task if the critic approves.
text_termination = TextMentionTermination("APPROVE")

# Create a team with the primary and critic agents.
team = RoundRobinGroupChat([primary_agent, critic_agent], termination_condition=text_termination)
# Use `asyncio.run(...)` when running in a script.
async def main():
    #result = await team.run(task="Write a short poem about the fall season.")
    #print(result)
    #await team.reset()  # Reset the team for a new task.
    await Console(team.run_stream(task="Write a short poem about the fall season."))  # Stream the messages to the console.
asyncio.run(main())